""" 
@author Erik M Volz
@date Aug 16 2011

Compute a simple likelihood profile for each birth rate in a randomly 
generated model. 
"""


from pylab import *
from scipy.interpolate import InterpolatedUnivariateSpline
from scipy import polyfit
from numpy import poly1d
import  pdb, csv, numpy,  time, sys , networkx
import modelGenerator, coalescentSimulator, likelihood


def bias_coverage(nprofiles=10, **kwargs):
	for n in range(nprofiles):
		try:
			f, liks_kl, birthRates,  genealogy =  likelihood_profile(**kwargs)
			
			ff = linspace(min(f), max(f), 10000)
			coverage = list() # true/false; covered the true value
			bias = list()
			for i, liks in enumerate(liks_kl):
				print 'birth rate', birthRates[i]
				fliks = array([ (f2, l2) for f2, l2 in zip(f, liks) if (l2 > -inf and l2 < inf)])
				#z = polyfit(fliks[:,0], fliks[:,1], 3)
				#p = poly1d(z.flatten())
				p = InterpolatedUnivariateSpline(fliks[:,0], fliks[:,1])
				lp = p(ff)
				i_maxLP = argmax(lp)
				if i_maxLP > 0:
					fhat = ff[argmax(lp)]
					flb = ff[  argmin( abs(max(lp) - 2 - lp[:argmax(lp)] )) ]
					fub =  ff[ argmax(lp) +  argmin( abs(max(lp) - 2 - lp[argmax(lp):] )) ]
				else: 
					fhat = ff[argmax(lp)]
					flb = ff[argmax(lp)]
					fub = ff[ argmax(lp) +  argmin( abs(max(lp) - 2 - lp[argmax(lp):] )) ]
				#
				
				bias.append(fhat - 1.)
				coverage.append(  (1>=flb and 1<=fub)  )
			#
		except:
			pass
	#
	
	try:
		return mean(bias), mean(coverage)
	except:
		bias,coverage
#
def likelihood_profile(m=6, n=700,  f = [.5, .75, .9, .95, 1., 1.05, 1.25, 1.5], 
		likelihoodClass = likelihood.Likelihood):
	""" 
	1. generates a model
	2. simulates coalescent tree
	3. for a single birth rate, calculates a likelihood profile by perturbing the rate by the factor f
	"""
	N=10**4
	n_births = m; n_migs = m; n_birth_loops = 0
	bg, mg = modelGenerator.generate_random_model__density(m, n_births, n_birth_loops, n_migs)
	ba = 2* array(networkx.adj_matrix(bg))
	for k in range(m):
		for l in range(m):
			ba[k,l] = ba[k,l] * (2 * uniform(.9, 1.1) ) #* uniform(.99, 1.01))
	ma = array(networkx.adj_matrix(mg))
	
	t = linspace(0., 15., 500)
	cm = modelGenerator.CompartmentalModel2(ba, ma, N, t)
	cm.births_migrations_prevalence()
	figure(-1)
	plot(cm.time, cm.Y.T)
	#show()
	
	sampleStates = ((cm.prevalence[-1])/(sum(cm.prevalence[-1])) * n).astype(int)
	k = n - sum(sampleStates)
	sampleStates[argmax(cm.prevalence[-1])] += k
	sampleStates = concatenate([[i]*ss for i,ss in enumerate(sampleStates)]).tolist()
	sampleStates = [ eye(m)[ss] for ss in sampleStates  ]
	sampleTimes = [max(cm.taxis)]*n #homochronous sample
	coalescentTree = coalescentSimulator.Simulator2( cm.taxis, cm.prevalence, cm.births, cm.migrations, sampleTimes, sampleStates)
	
	ssdict = dict( [  ('%i_' % i,  ss.tolist()) for i,ss in enumerate(sampleStates)] )  
	stdict = dict([  ('%i_' % i,st) for i,st in enumerate(sampleTimes) ])
	genealogy = likelihood.Genealogy(coalescentTree.newick, stdict, ssdict)
	
	kl_list = list()
	for k,l in [(k,l) for k in range(m) for l in range(m)]:
		if ba[k,l] > 0:
			kl_list.append( (k,l) )
	#
	#~ print k,l, ba
	
	ll1 = likelihoodClass(cm.taxis, cm.prevalence, cm.births, cm.migrations, genealogy).log_likelihood 
	
	birthRates = list()
	liks_kl = list()
	for k,l in kl_list:
		liks = list()
		birthRates.append( ba[k,l] )
		print 'calculating ', len(f), 'likelihoods', time.ctime()
		for ff in f:
			if ff==1.:
				liks.append(ll1)
			else:
				ba_f = copy(ba)
				ba_f[k,l] = ba_f[k,l] * ff
				cm_f = modelGenerator.CompartmentalModel2(ba_f, ma, N, t)
				cm_f.births_migrations_prevalence()
				liks.append( likelihoodClass(cm_f.taxis, cm_f.prevalence, cm_f.births, cm_f.migrations, genealogy).log_likelihood  )
				print 'calculated likelihood', ff, liks[-1], time.ctime()
		#
		liks_kl.append(liks)
	
	return f, liks_kl, birthRates, genealogy #f, liks, cm, z#f, f2, liks, liks2, cm


if __name__=='__main__':
	
	b, c = bias_coverage(nprofiles=100)
	#~ b,c = bias_coverage(nprofiles=20, likelihoodClass = likelihood.Likelihood)
	
	print 
	print 'bias, coverage'
	print b,c
	print 
	
	show()

